package cn.kevinking.ai.mcp.provider.configure;

import cn.kevinking.ai.mcp.provider.annotation.McpTool;
import cn.kevinking.ai.mcp.provider.annotation.McpToolScan;
import cn.kevinking.ai.mcp.provider.ex.McpProviderDuplicateToolException;
import cn.kevinking.ai.mcp.provider.utils.McpToolUtils;
import cn.kevinking.ai.mcp.provider.model.NativeToolDefinition;
import org.springframework.beans.factory.BeanDefinitionStoreException;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 该处只是通过mcp工具的注解注册definition, 不负责实例化对象
 *
 * @author kevinking
 */
public class McpToolScannerRegistrar implements ImportBeanDefinitionRegistrar {
	private static final AtomicInteger COUNTER = new AtomicInteger(0);

	@Override
	public void registerBeanDefinitions(AnnotationMetadata metadata, BeanDefinitionRegistry registry) {
		Map<String, Object> scanAttr = metadata.getAnnotationAttributes(McpToolScan.class.getName());
		assert scanAttr != null;

		List<String> basePackages = new java.util.ArrayList<>(Arrays.stream((String[]) scanAttr.get("basePackages"))
                .filter(StringUtils::hasText)
                .toList());
		if (basePackages.isEmpty()) {
			basePackages.add(getDefaultBasePackage(metadata));
		}

		// 拥有@McpTool方法的类必须被Spring容器接管
		ClassPathScanningCandidateComponentProvider scanner = new ClassPathScanningCandidateComponentProvider(true);
		basePackages.stream()
				.flatMap(basePackage -> scanner.findCandidateComponents(basePackage).stream())
				.forEach(bd -> {
					try {
						Class<?> beanClass = Class.forName(bd.getBeanClassName());
						List<Method> toolMethods = Arrays.stream(ReflectionUtils.getDeclaredMethods(beanClass))
								.filter(method -> method.isAnnotationPresent(McpTool.class))
								.toList();
						if (toolMethods.isEmpty()) {
							return;
						}
						if (checkDuplicateToolMethod(toolMethods)) {
							toolMethods.forEach(method -> registerToolDefinitionBean(registry, beanClass, method));
							return;
						}
						throw new McpProviderDuplicateToolException();
					} catch (ClassNotFoundException e) {
						throw new BeanDefinitionStoreException("fail to find bean class when registering mcp tools!");
					}
				});
	}

	private static void registerToolDefinitionBean(BeanDefinitionRegistry registry, Class<?> beanClass, Method method) {
		BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(NativeToolDefinition.class);
		McpTool toolAnno = method.getAnnotation(McpTool.class);
		builder.addPropertyValue("url", McpToolUtils.generateToolUrl(method));
		builder.addPropertyValue("description", toolAnno.description());
		// 这里不应该考虑去处paramDefinitions是否为空的情况，因为可能是无参数方法
		builder.addPropertyValue("paramDefinitions", McpToolUtils.from(method));
		builder.addPropertyValue("toolMethod", method);
		builder.addPropertyValue("beanClass", beanClass);
		registry.registerBeanDefinition(generateBaseBeanName(method.getName()), builder.getBeanDefinition());
	}

	private static String getDefaultBasePackage(AnnotationMetadata importingClassMetadata) {
		return ClassUtils.getPackageName(importingClassMetadata.getClassName());
	}

	/**
	 * 检查是否存在重复的工具方法
	 *
	 * @param methods 目标工具方法
	 * @return 重复判断结果, true为通过, false为不通过
	 */
	private static boolean checkDuplicateToolMethod(List<Method> methods) {
		Set<String> urlSet = methods.stream().map(McpToolUtils::generateToolUrl).collect(Collectors.toSet());
		return methods.size() == urlSet.size();
	}

	private static String generateBaseBeanName(String toolName) {
		 return String.format("%s#%s#%d", toolName, NativeToolDefinition.class.getSimpleName(), COUNTER.getAndIncrement());
	}
}
